/*************************************************************************
*  Wiizard - A Wii games manager
*  Copyright (C) 2023  CYBERDEViL
*
*  This file is part of Wiizard.
*
*  Wiizard is free software: you can redistribute it and/or modify
*  it under the terms of the GNU General Public License as published by
*  the Free Software Foundation, either version 3 of the License, or
*  (at your option) any later version.
*
*  Wiizard is distributed in the hope that it will be useful,
*  but WITHOUT ANY WARRANTY; without even the implied warranty of
*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
*  GNU General Public License for more details.
*
*  You should have received a copy of the GNU General Public License
*  along with this program.  If not, see <https://www.gnu.org/licenses/>.
*
*************************************************************************/

#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <unistd.h> // access
#include <libgen.h> // basename
#include <ctype.h>  // isdigit
#include <string.h> // strlen
#include <blkid/blkid.h>
#include <libfdisk/libfdisk.h>

// https://cdn.kernel.org/pub/linux/utils/util-linux/v2.29/libblkid-docs/libblkid-Low-level-probing.html
// https://cdn.kernel.org/pub/linux/utils/util-linux/v2.34/libfdisk-docs/libfdisk-Partition.html


static PyObject * PyBlockInfoError;


// strip tailing characters from a string
static void rstrip(char * text, char toStrip) {
  size_t strLen = strlen(text);
  char lastChar = text[strLen - 1];
  while (lastChar == ' ' && strLen > 0) {
    text[strLen - 1] = '\0';
    strLen   = strlen(text);
    lastChar = text[strLen - 1];
  }
}


static PyObject * method_getVendorModel(PyObject * self, PyObject * args) {
  // Allowed input:
  // - /dev/sdX
  // - /dev/sdXN
  // - sdX
  // - sdXN
  int errPy;
  char vendorPath[38]; // /sys/class/block/%s/device/vendor
  char modelPath[37];  // /sys/class/block/%s/device/model
  char * bname;
  char bnameBuff[6];
  char lastChar;
  size_t strLen;
  FILE * f;
  char vendorString[255];
  char modelString[255];
  char * devicePath = NULL;

  // parse the 'devicePath' argument
  errPy = PyArg_ParseTuple(args, "s", &devicePath);
  if (!errPy) {
    PyErr_SetString(PyBlockInfoError, "Invalid arguments");
    return NULL;
  }

  // get the base name in case devicePath is a /dev/sdX(N) path.
  bname = basename(devicePath); // do not free

  // do not allow basenames with a string length greater then 5 to prevent a
  // buffer overflow.
  if (strlen(bname) > 5) {
    PyErr_SetString(PyBlockInfoError, "Length of the basename is greater then 5");
    return NULL;
  }

  // copy the basename to a buff, we should not change the pointer we got from
  // basename.
  strncpy(bnameBuff, bname, sizeof(char) * 5);

  // create the vendor path
  if (snprintf(vendorPath, sizeof(char) * 37, "/sys/class/block/%s/device/vendor", bnameBuff) <= 0) {
    PyErr_SetString(PyBlockInfoError, "Failed to construct vendor path");
    return NULL;
  }

  // vendor path not found
  if (access(vendorPath, F_OK) != 0) {
    // now we try to strip all tailing digits from the basename, construct a
    // new vendor path and test if it exists.
    // so for example 'sda1' becomes 'sda'.
    strLen = strlen(bnameBuff);
    // strip all tailing digits
    for (size_t i=0; i < (strLen - 2); ++i) {
      strLen   = strlen(bnameBuff);
      lastChar = bnameBuff[strLen - 1];
      // current last character is a digit, strip it.
      if (isdigit(lastChar)) {
        bnameBuff[strLen - 1] = '\0';
      }
      // current last character is not a digit
      else {
        break;
      }
    }

    // construct a new vendor path from the tailing digit stripped basename.
    if (snprintf(vendorPath, sizeof(char) * 37, "/sys/class/block/%s/device/vendor", bnameBuff) <= 0) {
      PyErr_SetString(PyBlockInfoError, "Failed to construct vendor path (2)");
      return NULL;
    }
  }

  // construct the model path
  if (snprintf(modelPath, sizeof(char) * 36, "/sys/class/block/%s/device/model", bnameBuff) <= 0) {
    PyErr_SetString(PyBlockInfoError, "Failed to construct model path");
    return NULL;
  }

  // set as empty strings
  vendorString[0] = '\0';
  modelString[0]  = '\0';

  // Get the vendor string
  if (access(vendorPath, F_OK) == 0) {
    f = fopen(vendorPath, "r");
    if (f != NULL) {
      fgets(vendorString, 255, f);
      fclose(f);
    }
  }

  // Get the model string
  if (access(modelPath, F_OK) == 0) {
    f = fopen(modelPath, "r");
    if (f != NULL) {
      fgets(modelString, 255, f);
      fclose(f);
    }
  }

  // strip tailing new line
  vendorString[strcspn(vendorString, "\n")] = '\0';
  modelString[strcspn(modelString, "\n")] = '\0';

  // strip tailing spaces
  rstrip(vendorString, ' ');
  rstrip(modelString , ' ');

  //Py_RETURN_NONE;
  return Py_BuildValue("(s,s)", vendorString, modelString);
}


static PyObject * method_getBlkidValues(PyObject * self, PyObject * args) {
  int errPy;
  blkid_probe pr;
  const char * devicePath = NULL;

  // parse the 'devicePath' argument
  errPy = PyArg_ParseTuple(args, "s", &devicePath);
  if (!errPy) {
    PyErr_SetString(PyBlockInfoError, "Invalid arguments");
    return NULL;
  }

  // probe the device
  pr = blkid_new_probe_from_filename(devicePath);

  // failed to probe the device
  if (!pr) {
    PyErr_SetString(PyBlockInfoError, "Failed to create a new blkid probe");
    return NULL;
  }

  // actual probe
  if (blkid_do_probe(pr) != 0) {
    blkid_free_probe(pr);
    printf("Failed to probe '%s'\n", devicePath);
    PyErr_SetString(PyBlockInfoError, "Failed to probe");
    return NULL;
  }

  // create new dict object (this will we our result)
  PyObject * result = PyDict_New();

  // pointers for 'blkid_probe_get_value' to point their values to.
  const char * name;
  const char * data;
  size_t dataLen;

  // go over all values
  int nvals = blkid_probe_numof_values(pr);
  for (int n = 0; n < nvals; n++) {
    // get value
    if (blkid_probe_get_value(pr, n, &name, &data, &dataLen) == 0) {
      // No data (The data strings are \x00 terminated)
      if (dataLen <= 1) {
        continue;
      }
      // add to the result dict
      PyObject * dataStr = PyUnicode_FromStringAndSize(data, dataLen - 1);
      PyDict_SetItemString(result, name, dataStr);
    }
  }

  // cleanup
  blkid_free_probe(pr);

  return result;
}


static PyObject * method_getSize(PyObject * self, PyObject * args) {
  int err;
  const char * devicePath = NULL;
  struct fdisk_context * cxt;
  uint64_t bytes;

  // parse the 'devicePath' argument
  err = PyArg_ParseTuple(args, "s", &devicePath);
  if (!err) {
    PyErr_SetString(PyBlockInfoError, "Invalid arguments");
    return NULL;
  }

  cxt = fdisk_new_context();
  if (!cxt) {
    PyErr_SetString(PyBlockInfoError, "Failed to create new context");
    return NULL;
  }

  err = fdisk_assign_device(cxt, devicePath, 1); // read-only
  if (err) {
    fdisk_unref_context(cxt);
    PyErr_SetString(PyBlockInfoError, "Failed to open partition for reading");
    return NULL;
  }

  bytes = fdisk_get_nsectors(cxt) * fdisk_get_sector_size(cxt);

  // cleanup
  fdisk_deassign_device(cxt, 0);
  fdisk_unref_context(cxt);

  return PyLong_FromUnsignedLong(bytes);
}


static PyMethodDef pyBlockInfoMethods[] = {
  {"getBlkidValues", method_getBlkidValues, METH_VARARGS,
    "Returns all values from libblkid for given device."},

  {"getVendorModel", method_getVendorModel, METH_VARARGS,
    "Get a block device its vendor and model strings."},

  {"getSize", method_getSize, METH_VARARGS,
    "Returns the size in bytes of a partition or disk."},

  {NULL, NULL, 0, NULL}
};


static struct PyModuleDef pyBlockInfoModule = {
  PyModuleDef_HEAD_INIT,
  "pyblockinfo",
  "Get extra info about block devices",
  -1,
  pyBlockInfoMethods
};


PyMODINIT_FUNC PyInit_pyblockinfo(void) {
  PyObject * module = PyModule_Create(&pyBlockInfoModule);

  // https://docs.python.org/3/c-api/exceptions.html#exception-classes
  PyBlockInfoError = PyErr_NewException("pyblockinfo.error", NULL, NULL);
  Py_XINCREF(PyBlockInfoError);
  if (PyModule_AddObject(module, "error", PyBlockInfoError) < 0) {
    Py_XDECREF(PyBlockInfoError);
    Py_CLEAR(PyBlockInfoError);
    Py_DECREF(module);
    return NULL;
  }

  return module;
}
